import pandas as pd
import datetime as dt

from account import Account


class BackTest:
    '''
    回测类：包含选股策略
    '''
    def __init__(self, cfg, stocklist):#初始化回测对象
        self.skid_rate = cfg['skid_rate']
        self.fee_rate = cfg['fee_rate']
        self.end_date = dt.datetime.strptime(cfg['end_date'], '%Y-%m-%d')
        self.account = Account(cfg['cash'], self.skid_rate, self.fee_rate)
        self.stocklist = stocklist

        self.data = None # 今天的数据
        self.last_data = None # 前一天的数据
        self.date = None
        self.last_date = None

        self.buy_orders = {}
        self.sell_orders = {}

        self.cum_chginday = pd.Series([0]*(len(self.stocklist)), index=stocklist)#统计累计日内变化

    def update_data(self, data, date):#更新数据，把新的一天数据传出来
        self.last_data = self.data
        self.data = data
        self.last_date = self.date
        self.date = date

    def check_criterion(self):#判断今天要进行买的和卖的股票
        '''
        交易策略：选择前一天跌的最多的十支股票买入
        :return:
        '''
        if self.last_date:
            self.cum_chginday += (self.last_data['close'] / self.last_data['open'] - 1)#计算前一天收益率
        else:
            return

        if self.date.isocalendar()[2] == 1 and self.date != self.end_date:#只在周一交易
            buy_list = self.cum_chginday.nsmallest(10).index.tolist()#找出跌最多的10
            self.account.update_equity(self.data['open'].to_dict()) # 股票价格变动导致资产价值变化
            cashtobuy = round(self.account.equity / (1+self.skid_rate) ** 2 / (1+self.fee_rate) ** 2 / 10, 2)#计算今天能有多少钱来买一支股票 钱平均分给10股票
            for code in buy_list:
                price = self.data.loc[code, 'open']
                self.buy_orders[code] = [price, int(cashtobuy / price / 100) * 100, self.date.strftime('%Y-%m-%d')]
            for code in self.account.position:
                price = self.data.loc[code, 'open']
                self.sell_orders[code] = [price, self.account.position[code], self.date.strftime('%Y-%m-%d')]
        elif self.date == self.end_date:#最后一天清仓
            for code in self.account.position:
                price = self.data.loc[code, 'close']
                self.sell_orders[code] = [price, self.account.position[code], self.date.strftime('%Y-%m-%d')]
        else:
            pass
        return

    def trigger_strategy(self):#过去一周日内跌最多的10股票票买入
        '''
        选择跌的最多的是指股票
        :return:
        '''
        self.check_criterion()

        for code in self.sell_orders:#对要卖的股票检查重合部分 持有的也有可能是跌的最多的10支股票
            if code in self.buy_orders:
                if self.sell_orders[code][1] > self.buy_orders[code][1]:
                    self.sell_orders[code][1] -= self.buy_orders[code][1]
                    self.buy_orders[code][1] = 0
                    self.account.sell_order(code, *self.sell_orders[code])
                elif self.sell_orders[code][1] < self.buy_orders[code][1]:
                    self.buy_orders[code][1] -= self.sell_orders[code][1]
                else:
                    pass
            else:
                self.account.sell_order(code, *self.sell_orders[code])

        for code in self.buy_orders:
            self.account.buy_order(code, *self.buy_orders[code])

        self.buy_orders = {}
        self.sell_orders = {}

        self.account.update_equity(self.data['close'].to_dict())
        self.account.update_wallet(self.date.strftime('%Y-%m-%d'))